#! /usr/bin/env python
#
# couchdb-xapian-indexer - Index couchdb databases with Xapian.
# Copyright (C) 2008 Paul J. Davis <paul.joseph.davis@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import logging
import os
import sys
from optparse import OptionParser, make_option

import simplejson
import xapian

log = logging.getLogger(__name__)

class QueryHandler(object):
    def __init__(self, dir, url, dbname):
        self.idxfile = os.path.join(dir, '%s.idx' % dbname)
        self.db = xapian.Database(self.idxfile)
        self.enq = xapian.Enquire(self.db)
        self.qp = xapian.QueryParser()
        self.qp.set_stemmer(xapian.Stem("english"))
        self.qp.set_database(self.db)
        self.qp.set_stemming_strategy(xapian.QueryParser.STEM_SOME)
    def query(self, query):
        self.db.reopen()
        qs = query.get('query', {}).get('q', '')
        offset = query.get('query', {}).get('offset', 0)
        count = query.get('query', {}).get('count', 25)
        q = self.qp.parse_query(qs)
        self.enq.set_query(q)
        ret = []
        mset = self.enq.get_mset(offset, count)
        log.info(dir(mset))
        for m in mset:
            ret.append({'id': m.document.get_data(), 'rank': m.rank + 1, 'score': m.percent})
        return {'docs': ret, 'count': mset.get_matches_estimated()}

def queries():
    line = sys.stdin.readline()
    while line:
        if not line:
            return
        obj = simplejson.loads(line)
        yield obj
        line = sys.stdin.readline()

def send(data):
    sys.stdout.write(simplejson.dumps(data))
    sys.stdout.write('\n')
    sys.stdout.flush()

def main(dir, url, exclude):
    if not os.path.isdir(dir):
        os.mkdir(dir)
    handlers = {}
    for query in queries():
        if query == [u'reset']:
            send(True)
            continue
        log.info("Update: %s" % query)
        dbname = query['db']
        if dbname in exclude:
            send({'code': '404', 'json': {'error': 'Database not indexed.'}, 'headers': {}})
            continue
        if dbname not in handlers:
            handlers[dbname] = QueryHandler(dir, url, dbname)
        ret = handlers[dbname].query(query)
        send({'code': '200', 'json': ret, 'headers': {}})

if __name__ == '__main__':
    options = [
        make_option('-d', '--dir', dest='dir', metavar="DIRECTORY", default="./xapian",
            help="Directory in which to store xapian databases. [%default]"),
        make_option('-u', '--url', dest='url', metavar="URL", default="http://localhost:5984",
            help="URL of the couchdb server. [%default]"),
        make_option('-e', '--exclude', dest='exclude', metavar='DB_NAME', default=[],
            help="Exclude a database from indexing. Can be used multiple times."),
        make_option('-l', '--log', dest='log', metavar="FILE", default='./xapian/query.log',
            help="Name of the log file to write to."),
    ]
    parser = OptionParser("usage: %prog [OPTIONS]", option_list=options)
    opts, args = parser.parse_args()
    if len(args) != 0:
        print "Unrecognized arguments: %s" % ' '.join(args)
        parser.print_help()
        exit(-1)
    logging.basicConfig(filename=opts.log, level=logging.DEBUG, format="%(levelname)s %(message)s")
    try:
        main(os.path.abspath(opts.dir), opts.url, opts.exclude)
    except:
        log.exception("Querying shutting down due to high stress. Relaxation needed.")
